from datasets import load_dataset, Dataset, Audio
import os

# --- 步骤 1: 加载原始的 TSV 文件 ---
# 假设你的 TSV 文件就在 nl 文件夹里
tsv_file_path = "./voxpopuli-data/nl/asr_train.tsv"  # <-- 请替换成你的实际 TSV 文件名

# 使用通用的 'csv' 加载器，并指定分隔符为 '\t'
raw_dataset = load_dataset('csv', data_files=tsv_file_path, delimiter='\t', split='train')
raw_dataset = raw_dataset.select(range(100))
# 检查初始加载的列名
print("原始 TSV 加载后的列名:")
print(raw_dataset.column_names)
# 预期输出: ['audio_id', 'raw_text', 'normalized_text', 'speaker_id', 'gender', 'is_gold_transcript', 'accent']

# --- 步骤 2: 添加 'language' 列 ---
# 为所有行添加一个固定值 "nl" 的新列
dataset_with_lang = raw_dataset.add_column('language', ['nl'] * len(raw_dataset))

print("\n添加 'language' 列后的列名:")
print(dataset_with_lang.column_names)

# --- 步骤 3: 构建 'audio' 列 ---
# 这是核心步骤。我们需要一个函数来根据 audio_id 生成完整的文件路径。

# 首先，定义你的音频文件所在的根目录
# 这个目录应该包含 "2014", "2015" 等年份子文件夹
audio_root = "./voxpopuli-data/nl/train/train_part_0/"  # <-- 请确认你的音频文件实际存放位置


def build_audio_path(example):
    """
    这个函数接收数据集中的一行数据（一个 example 字典），
    并返回一个包含新 'audio' 列值的字典。
    """
    audio_id = example['id']



    # 使用 os.path.join 来构建路径，这能确保在不同操作系统下都有效
    full_path = os.path.join(audio_root, f"{audio_id}.wav")

    # 返回一个新的字典，键是我们要创建的列名 'audio'，值是完整的路径
    return {'audio': full_path}


# 使用 map 函数将这个构建逻辑应用到数据集的每一行
# 这会创建一个名为 'audio' 的新列，其值为文件路径字符串
dataset_with_path = dataset_with_lang.map(build_audio_path)

print("\n添加 'audio' 路径列后的列名:")
print(dataset_with_path.column_names)
# 现在你可以看到 'audio' 列已经存在了，但它还是字符串类型
print("\n'audio' 列的第一个元素（目前是路径字符串）:")
print(dataset_with_path[0]['audio'])
s=dataset_with_path[0]
print(s)
# --- 步骤 4: 将 'audio' 列转换为真正的 Audio 对象 ---
# 使用 cast_column 告诉 datasets 库，'audio' 列应该被当作音频文件来处理
final_dataset = dataset_with_path.cast_column('audio', Audio(sampling_rate=16000))

print("\n使用 cast_column 后的列名:")
print(final_dataset.column_names)
# 列名没变，但列的类型已经变了

print("\n'audio' 列的第一个元素（现在是 Audio 对象）:")
print(final_dataset[0]['audio'])
print('final_dataset[0]:',final_dataset[0])
# 预期输出类似: {'path': '...', 'array': array([...], dtype=int16), 'sampling_rate': 16000}


# # --- 步骤 5: 验证并继续你的后续操作 ---
# # 现在，你的 final_dataset 结构就完全正确了！
# print("\n最终数据集的列名（与标准 VoxPopuli 一致）:")
# print(final_dataset.column_names)
# # 预期输出: ['audio_id', 'raw_text', 'normalized_text', 'speaker_id', 'gender', 'is_gold_transcript', 'accent', 'language', 'audio']
#
# # 选择前 100 条数据进行后续处理
# dataset_to_process = final_dataset.select(range(100))


#["audio_id", "raw_text", "normalized_text", "speaker_id", "gender","is_gold_transcript","accent"]
#Column names: ['audio_id', 'language', 'audio', 'raw_text', 'normalized_text', 'gender', 'speaker_id', 'is_gold_transcript', 'accent']

import os
import json

# --- 1. 配置 ---
# 设置你要遍历的根目录
# 请将此路径修改为你的音频文件实际存放位置
root_directory = "./voxpopuli-data/nl/train/"

# 设置输出的 JSON 文件名
output_json_file = "audio_files_index.json"

# --- 2. 遍历目录并收集信息 ---
# 创建一个空列表，用于存储所有文件的信息
audio_files_data = {}

print(f"开始遍历目录: {root_directory}")

# os.walk 会生成目录树中的每个目录的元组 (dirpath, dirnames, filenames)
for dirpath, dirnames, filenames in os.walk(root_directory):
    for filename in filenames:
        # 我们只关心 .wav 文件，可以根据需要修改扩展名
        if filename.lower().endswith(".wav"):
            # 获取文件的完整路径
            full_path = os.path.join(dirpath, filename)

            # 获取文件大小（字节）
            file_size = os.path.getsize(full_path)

            # 创建一个字典来存储当前文件的信息
            file_info = {
                "filename": filename,
                "path": full_path,
                "size_bytes": file_size,
                # 你可以在这里添加更多信息，比如从路径中提取的元数据
                # "parent_folder": os.path.basename(dirpath)
            }

            # 将这个文件的信息字典添加到我们的列表中
            audio_files_data[filename.replace('.wav','')]=file_info

print(f"遍历完成，共找到 {len(audio_files_data)} 个音频文件。")

# --- 3. 将结果保存为 JSON 文件 ---
print(f"正在将结果保存到 {output_json_file} ...")

# 使用 json.dump 将列表写入文件
# indent=4 使 JSON 文件格式化，易于阅读
# ensure_ascii=False 确保非英文字符（如中文）能正确保存
try:
    with open(output_json_file, 'w', encoding='utf-8') as f:
        json.dump(audio_files_data, f, indent=4, ensure_ascii=False)
    print("保存成功！")
except IOError as e:
    print(f"保存文件时出错: {e}")

